import os

import numpy as np
from tqdm import tqdm

def generate_doa_dataset(num_data, num_sensors, max_signals_per_sample, 
                         snr_db, min_angle=0, max_angle=np.pi, 
                         signal_length=1000, sensor_distance=0.1, 
                         wavelength=0.343):
  """
  生成 DOA 波达方位估计仿真数据集

  Args:
    num_data: 样本数量
    num_sensors: 传感器数量
    max_signals_per_sample: 每个样本中最大信号源数量
    snr_db: 信噪比 (dB)
    min_angle: 角度范围最小值 (弧度)
    max_angle: 角度范围最大值 (弧度)
    signal_length: 信号长度 (采样点数)
    sensor_distance: 传感器间距
    wavelength: 信号波长

  Returns:
    signals: 仿真信号数据集，形状为 (num_data, signal_length, num_sensors)
    labels: 标签数据集，列表，每个元素为一个样本对应的角度真值列表
  """

  signals = []
  labels = []

  for _ in tqdm(range(num_data)):
    # 随机生成信号源数量
    num_signals = np.random.randint(1, max_signals_per_sample + 1)

    # 随机生成信号源角度
    signal_angles = np.random.uniform(min_angle, max_angle, num_signals)

    # 生成仿真信号
    array_signal = np.zeros((signal_length, num_sensors), dtype=complex)
    d_over_lambda = sensor_distance / wavelength

    # 生成正弦波时间序列
    time = np.arange(signal_length) / signal_length
    source_signal = np.exp(1j * 2 * np.pi * time)[:,None]

    for angle in signal_angles:
        # 计算每个传感器的相位延迟
        phase_shifts = 2 * np.pi * np.arange(num_sensors) * np.cos(angle) * d_over_lambda

        # 应用相位延迟到正弦波
        sensor_signals = source_signal * np.exp(1j * phase_shifts)

        # 将传感器信号叠加到阵列信号
        array_signal += sensor_signals
    # 添加噪声
    noise_power = 10**(-snr_db / 10)
    noise = np.random.normal(0, np.sqrt(noise_power / 2), size=sensor_signals.shape) + \
            1j * np.random.normal(0, np.sqrt(noise_power / 2), size=sensor_signals.shape)
    sensor_signals += noise
    # 将信号转换为实数并展平
    array_signal = np.concatenate((np.real(array_signal), np.imag(array_signal)), axis=1)
    # array_signal = array_signal.flatten()

    # 添加到数据集
    signals.append(array_signal)
    labels.append(signal_angles.tolist())

  # 转换为 numpy 数组
  signals = np.array(signals)

  return signals, labels



k_data = 1
num_data = 1000*k_data
num_sensors = 6
SNR = 3
# 生成数据集
signals, labels = generate_doa_dataset(num_data=num_data, num_sensors=num_sensors, max_signals_per_sample=3, snr_db=SNR)

# 保存数据集
os.makedirs("data", exist_ok=True)
np.save(f"data/signals_{k_data}k_data_{num_sensors}sensors_{SNR}SNR.npy", signals)
np.save(f"data/labels_{k_data}k_data_{num_sensors}sensors_{SNR}SNR.npy", labels)